import time
from random import SystemRandom
from sm3 import *
import func

# 选择素域，设置椭圆曲线参数
default_ecc_table = {
    'n': 'FFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFF7203DF6B21C6052B53BBF40939D54123',
    'p': 'FFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00000000FFFFFFFFFFFFFFFF',
    'g': '32c4ae2c1f1981195f9904466a39c9948fe30bbff2660be1715a4589334c74c7'
         'bc3736a2f4f6779c59bdcee36b692153d0a9877cc62a474002df32e52139f0a0',
    'a': 'FFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00000000FFFFFFFFFFFFFFFC',
    'b': '28E9FA9E9D9F5E344D5A9E4BCF6509A7F39789F515AB8F92DDBCBD414D940E93',
}


# 拓展欧几里得用于计算逆元
def exgcd(a, b):
    if b == 0:
        return 1, 0, a
    else:
        x, y, q = exgcd(b, a % b)
        x, y = y, (x - (a // b) * y)
        return x, y, q


def inf(a, p):
    x, y, q = exgcd(a, p)
    if q != 1:
        raise Exception("No solution.")
    else:
        return (x + p) % p  # 防止负数


class CryptSM2(object):

    def __init__(self, ecc_table=default_ecc_table, mode=0, public_key=None, private_key=None):
        """
        mode: 0-C1C2C3, 1-C1C3C2 (default is 1)
        """
        self.para_len = len(ecc_table['n'])
        self.ecc_a3 = (
                              int(ecc_table['a'], base=16) + 3) % int(ecc_table['p'], base=16)
        self.ecc_table = ecc_table
        assert mode in (0, 1), 'mode must be one of (0, 1)'
        self.public_key = public_key
        self.private_key = private_key
        self.mode = mode
        self.d1 = None
        self.d2 = None
        self.p1 = None
        self.p2 = None
        self.P1 = None
        self.P2 = None
        self.data = b"1111"
        self.k1 = None
        self.k2 = None
        self.k3 = None
        self.Q1 = None
        self.Q2 = None
        self.r1 = None
        self.r2 = None
        self.c1 = None
        self.c2 = None
        self.s1 = None
        self.s2 = None
        self.RS = None

    def _kg(self, k, Point):  # kP运算
        Point = '%s%s' % (Point, '1')
        mask_str = '8'
        for i in range(self.para_len - 1):
            mask_str += '0'
        mask = int(mask_str, 16)
        Temp = Point
        flag = False
        for n in range(self.para_len * 4):
            if (flag):
                Temp = self._double_point(Temp)
            if (k & mask) != 0:
                if (flag):
                    Temp = self._add_point(Temp, Point)
                else:
                    flag = True
                    Temp = Point
            k = k << 1
        return self._convert_jacb_to_nor(Temp)

    def _double_point(self, Point):  # 倍点
        l = len(Point)
        len_2 = 2 * self.para_len
        if l < self.para_len * 2:
            return None
        else:
            x1 = int(Point[0:self.para_len], 16)
            y1 = int(Point[self.para_len:len_2], 16)
            if l == len_2:
                z1 = 1
            else:
                z1 = int(Point[len_2:], 16)

            T6 = (z1 * z1) % int(self.ecc_table['p'], base=16)
            T2 = (y1 * y1) % int(self.ecc_table['p'], base=16)
            T3 = (x1 + T6) % int(self.ecc_table['p'], base=16)
            T4 = (x1 - T6) % int(self.ecc_table['p'], base=16)
            T1 = (T3 * T4) % int(self.ecc_table['p'], base=16)
            T3 = (y1 * z1) % int(self.ecc_table['p'], base=16)
            T4 = (T2 * 8) % int(self.ecc_table['p'], base=16)
            T5 = (x1 * T4) % int(self.ecc_table['p'], base=16)
            T1 = (T1 * 3) % int(self.ecc_table['p'], base=16)
            T6 = (T6 * T6) % int(self.ecc_table['p'], base=16)
            T6 = (self.ecc_a3 * T6) % int(self.ecc_table['p'], base=16)
            T1 = (T1 + T6) % int(self.ecc_table['p'], base=16)
            z3 = (T3 + T3) % int(self.ecc_table['p'], base=16)
            T3 = (T1 * T1) % int(self.ecc_table['p'], base=16)
            T2 = (T2 * T4) % int(self.ecc_table['p'], base=16)
            x3 = (T3 - T5) % int(self.ecc_table['p'], base=16)

            if (T5 % 2) == 1:
                T4 = (T5 + ((T5 + int(self.ecc_table['p'], base=16)) >> 1) - T3) % int(
                    self.ecc_table['p'], base=16)
            else:
                T4 = (T5 + (T5 >> 1) - T3) % int(self.ecc_table['p'], base=16)

            T1 = (T1 * T4) % int(self.ecc_table['p'], base=16)
            y3 = (T1 - T2) % int(self.ecc_table['p'], base=16)

            form = '%%0%dx' % self.para_len
            form = form * 3
            return form % (x3, y3, z3)

    def _add_point(self, P1, P2):  # 点加函数，P2点为仿射坐标即z=1，P1为Jacobian加重射影坐标
        len_2 = 2 * self.para_len
        l1 = len(P1)
        l2 = len(P2)
        if (l1 < len_2) or (l2 < len_2):
            return None
        else:
            X1 = int(P1[0:self.para_len], 16)
            Y1 = int(P1[self.para_len:len_2], 16)
            if (l1 == len_2):
                Z1 = 1
            else:
                Z1 = int(P1[len_2:], 16)
            x2 = int(P2[0:self.para_len], 16)
            y2 = int(P2[self.para_len:len_2], 16)

            T1 = (Z1 * Z1) % int(self.ecc_table['p'], base=16)
            T2 = (y2 * Z1) % int(self.ecc_table['p'], base=16)
            T3 = (x2 * T1) % int(self.ecc_table['p'], base=16)
            T1 = (T1 * T2) % int(self.ecc_table['p'], base=16)
            T2 = (T3 - X1) % int(self.ecc_table['p'], base=16)
            T3 = (T3 + X1) % int(self.ecc_table['p'], base=16)
            T4 = (T2 * T2) % int(self.ecc_table['p'], base=16)
            T1 = (T1 - Y1) % int(self.ecc_table['p'], base=16)
            Z3 = (Z1 * T2) % int(self.ecc_table['p'], base=16)
            T2 = (T2 * T4) % int(self.ecc_table['p'], base=16)
            T3 = (T3 * T4) % int(self.ecc_table['p'], base=16)
            T5 = (T1 * T1) % int(self.ecc_table['p'], base=16)
            T4 = (X1 * T4) % int(self.ecc_table['p'], base=16)
            X3 = (T5 - T3) % int(self.ecc_table['p'], base=16)
            T2 = (Y1 * T2) % int(self.ecc_table['p'], base=16)
            T3 = (T4 - X3) % int(self.ecc_table['p'], base=16)
            T1 = (T1 * T3) % int(self.ecc_table['p'], base=16)
            Y3 = (T1 - T2) % int(self.ecc_table['p'], base=16)

            form = '%%0%dx' % self.para_len
            form = form * 3
            return form % (X3, Y3, Z3)

    def _minus_point(self, P1, P2):  # 点减函数，P2点为仿射坐标即z=1，P1为Jacobian加重射影坐标
        len_2 = 2 * self.para_len
        l1 = len(P1)
        l2 = len(P2)
        if (l1 < len_2) or (l2 < len_2):
            return None
        else:
            X1 = int(P1[0:self.para_len], 16)
            Y1 = int(P1[self.para_len:len_2], 16)
            if (l1 == len_2):
                Z1 = 1
            else:
                Z1 = int(P1[len_2:], 16)
            x2 = int(P2[0:self.para_len], 16)
            y2 = int(P2[self.para_len:len_2], 16)
            y2 = -y2  # 将第二个点的y轴坐标变为负值

            T1 = (Z1 * Z1) % int(self.ecc_table['p'], base=16)
            T2 = (y2 * Z1) % int(self.ecc_table['p'], base=16)
            T3 = (x2 * T1) % int(self.ecc_table['p'], base=16)
            T1 = (T1 * T2) % int(self.ecc_table['p'], base=16)
            T2 = (T3 - X1) % int(self.ecc_table['p'], base=16)
            T3 = (T3 + X1) % int(self.ecc_table['p'], base=16)
            T4 = (T2 * T2) % int(self.ecc_table['p'], base=16)
            T1 = (T1 - Y1) % int(self.ecc_table['p'], base=16)
            Z3 = (Z1 * T2) % int(self.ecc_table['p'], base=16)
            T2 = (T2 * T4) % int(self.ecc_table['p'], base=16)
            T3 = (T3 * T4) % int(self.ecc_table['p'], base=16)
            T5 = (T1 * T1) % int(self.ecc_table['p'], base=16)
            T4 = (X1 * T4) % int(self.ecc_table['p'], base=16)
            X3 = (T5 - T3) % int(self.ecc_table['p'], base=16)
            T2 = (Y1 * T2) % int(self.ecc_table['p'], base=16)
            T3 = (T4 - X3) % int(self.ecc_table['p'], base=16)
            T1 = (T1 * T3) % int(self.ecc_table['p'], base=16)
            Y3 = (T1 - T2) % int(self.ecc_table['p'], base=16)

            form = '%%0%dx' % self.para_len
            form = form * 3
            return form % (X3, Y3, Z3)

    def _convert_jacb_to_nor(self, Point):  # Jacobian加重射影坐标转换成仿射坐标
        len_2 = 2 * self.para_len
        x = int(Point[0:self.para_len], 16)
        y = int(Point[self.para_len:len_2], 16)
        z = int(Point[len_2:], 16)
        z_inv = pow(
            z, int(self.ecc_table['p'], base=16) - 2, int(self.ecc_table['p'], base=16))
        z_invSquar = (z_inv * z_inv) % int(self.ecc_table['p'], base=16)
        z_invQube = (z_invSquar * z_inv) % int(self.ecc_table['p'], base=16)
        x_new = (x * z_invSquar) % int(self.ecc_table['p'], base=16)
        y_new = (y * z_invQube) % int(self.ecc_table['p'], base=16)
        z_new = (z * z_inv) % int(self.ecc_table['p'], base=16)
        if z_new == 1:
            form = '%%0%dx' % self.para_len
            form = form * 2
            return form % (x_new, y_new)
        else:
            return None

    def verify(self, Sign, data):
        # 验签函数，sign签名r||s，E消息hash，public_key公钥
        r = int(Sign[0:self.para_len], 16)
        s = int(Sign[self.para_len:2 * self.para_len], 16)
        e = int(data.hex(), 16)
        t = (r + s) % int(self.ecc_table['n'], base=16)
        if t == 0:
            return 0

        P1 = self._kg(s, self.ecc_table['g'])
        P2 = self._kg(t, self.public_key)
        # print(P1)
        # print(P2)
        if P1 == P2:
            P1 = '%s%s' % (P1, 1)
            P1 = self._double_point(P1)
        else:
            P1 = '%s%s' % (P1, 1)
            # print(P1)
            # print(P2)
            P1 = self._add_point(P1, P2)
            P1 = self._convert_jacb_to_nor(P1)

        x = int(P1[0:self.para_len], 16)
        return (r == ((e + x) % int(self.ecc_table['n'], base=16)))

    def cilent_out_d1p1(self):
        self.d1 = SystemRandom().randrange(1, int(self.ecc_table['n'], 16) - 1)  # 输出: 十进制 int类型
        self.p1 = self._kg(self.d1, self.ecc_table['g'])  # 输出: 十六进制 128bit str类型

    def server_out_d2p2(self):
        self.d2 = SystemRandom().randrange(1, int(self.ecc_table['n'], 16) - 1)  # 输出: 十进制 int类型
        self.p2 = self._kg(self.d2, self.ecc_table['g'])  # 输出: 十六进制 128bit str类型

    def cilent_out_P1(self, p2):
        d1p2 = self._kg(self.d1, p2)  # 输出: 十六进制 128bit str类型
        # print(d1p2)
        P1 = '%s%s' % (d1p2, 1)  # 前提是d1p2和基点G不相等 否则建议使用二倍点 在数据的后面拼接一个1
        P1 = self._minus_point(P1, self.ecc_table['g'])  # 进行点减运算
        P1 = self._convert_jacb_to_nor(P1)
        self.P1 = P1

    def server_out_P2(self, p1):
        d2p1 = self._kg(self.d2, p1)  # 输出: 十六进制 128bit str类型
        # print(d2p1)
        P2 = '%s%s' % (d2p1, 1)  # 前提是d1p2和基点G不相等 否则建议使用二倍点
        # print(P1)
        # print(P2)
        P2 = self._minus_point(P2, self.ecc_table['g'])  # 进行点减运算
        P2 = self._convert_jacb_to_nor(P2)
        self.P2 = P2

    def cilent_out_k1Q1(self):
        self.k1 = SystemRandom().randrange(1, int(self.ecc_table['n'], 16))  # 输出: 十进制 int类型
        self.Q1 = self._kg(self.k1, self.ecc_table['g'])  # 输出: 十六进制 128bit str类型

    def server_out_k2Q2(self):
        self.k2 = SystemRandom().randrange(1, int(self.ecc_table['n'], 16))  # 输出: 十进制 int类型
        self.Q2 = self._kg(self.k2, self.ecc_table['g'])  # 输出: 十六进制 128bit str类型

    def server_out_r2c1c1(self, Q1):
        k2Q1 = self._kg(self.k2, Q1)  # 输出: 十六进制 128bit str类型
        E = self.data.hex()  # 消息转化为16进制字符串
        e = int(E, 16)  # 消息转化为十进制字符串
        x1 = int(k2Q1[0:self.para_len], 16)  # 截取前64位即x1 转换为十进制
        self.r2 = ((x1 + e) % int(self.ecc_table['n'], base=16))

        self.k3 = SystemRandom().randrange(1, int(self.ecc_table['n'], 16))  # 输出: 十进制 int类型
        d2_1 = inf(self.d2, int(self.ecc_table['n'], 16))  # 计算d2的逆元输出为十进制数
        self.c1 = ((d2_1 * self.k2 * self.k3) % int(self.ecc_table['n'], base=16))
        self.c2 = (d2_1 * self.r2 * (1 - self.k2)) % int(self.ecc_table['n'], base=16)

    def cilent_out_r1s1(self, Q2, c1):
        k1Q2 = self._kg(self.k1, Q2)  # 输出: 十六进制 128bit str类型
        E = self.data.hex()  # 消息转化为16进制字符串
        e = int(E, 16)  # 消息转化为十进制字符串
        x1 = int(k1Q2[0:self.para_len], 16)  # 截取前64位即x1 转换为十进制
        self.r1 = ((x1 + e) % int(self.ecc_table['n'], base=16))

        d1_1 = inf(self.d1, int(self.ecc_table['n'], 16))  # 计算d1的逆元输出为十进制数
        self.s1 = (d1_1 * c1 * (self.k1 + self.r1)) % int(self.ecc_table['n'], base=16)

    def server_out_s2(self, s1):
        k3_1 = inf(self.k3, int(self.ecc_table['n'], 16))  # 计算k3的逆元输出为十进制数
        self.s2 = (s1 * k3_1) % int(self.ecc_table['n'], base=16)

    def cilent_out_s(self, s2, c2):
        d1_1 = inf(self.d1, int(self.ecc_table['n'], 16))  # 计算d1的逆元输出为十进制数
        s = (s2 + (d1_1 * c2) - self.r1) % int(self.ecc_table['n'], base=16)
        self.RS = '%064x%064x' % (self.r1, s)
        return '%064x%064x' % (self.r1, s)

    def message_to_sm3_hash(self, message: list):
        string = ""
        for i in message:
            string = string + str(i)
        string = string.encode(encoding="utf-8")
        message_hash = sm3_hash(func.bytes_to_list(string))
        self.data = message_hash.encode(encoding="utf-8")
        return message_hash


if __name__ == '__main__':
    time1 = time.time()

    cilent = CryptSM2()
    server = CryptSM2()

    cilent.cilent_out_d1p1()
    server.server_out_d2p2()

    cilent.cilent_out_P1(server.p2)
    time2 = time.time()
    server.server_out_P2(cilent.p1)
    time3 = time.time()

    cilent.cilent_out_k1Q1()
    server.server_out_k2Q2()

    server.server_out_r2c1c1(Q1=cilent.Q1)
    cilent.cilent_out_r1s1(Q2=server.Q2, c1=server.c1)

    server.server_out_s2(s1=cilent.s1)
    time4 = time.time()
    cilent.cilent_out_s(s2=server.s2, c2=server.c2)
    time5 = time.time()

    print(f"服务端创建公钥用时：{int((time3 - time1) * 1000)}ms")
    print(f"客户端创建公钥用时：{int((time2 - time1) * 1000)}ms")
    print(f"服务端协作签名用时：{int((time4 - time3) * 1000)}ms")
    print(f"客户端协作签名用时：{int((time5 - time3) * 1000)}ms")

